import numpy as np

from core.cuda import cuda_module
from core.layer import Layer
from implement.variables.parameter import Parameter
from utils.functions_collect import linear


# 线性层
class Affine(Layer):
    def __init__(self, out_size, in_size=None, nobias=False, dtype=np.float32):
        """
        设置对象的属性。如果属性值为Parameter或Layer类的实例，将其添加到_params集合中。

        Parameters:
            out_size: 表示神经网络层的输出尺寸，即该层神经元的数量。
            in_size: 表示神经网络层的输入尺寸，即该层接收的输入数据的特征数量。这个参数是可选的，如果没有提供，则会在后续处理中进行初始化。
            nobias: 一个布尔值，表示是否使用偏置。如果为True，表示不使用偏置；如果为False，表示使用偏置。
            dtype: 指定权重和偏置的数据类型，默认为32位浮点数

        Returns:
            None
        """
        # 调用父类的构造函数
        super().__init__()

        # 延迟创建权重W的时间
        self.in_size = in_size
        self.out_size = out_size
        self.dtype = dtype
        self.xp = cuda_module  # 使用的是cuda模块，如果没有启用GPU加速，则使用numpy

        # 初始化权重参数W，但不立即赋值权重的具体数值
        self.W = Parameter(None, name='W')

        # 如果没有指定in_size则延后处理
        if self.in_size is not None:
            self._init_W()

        # 如果设置nobias为True，则不使用偏置
        if nobias:
            self.b = None
        else:
            # 初始化偏置参数b，全零初始化
            self.b = Parameter(self.xp.zeros(out_size, dtype=dtype), name='b')

    def _init_W(self):
        I, O = self.in_size, self.out_size
        W_data = self.xp.random.randn(I, O).astype(self.dtype) * self.xp.sqrt(1 / I)
        self.W.data = W_data

    def forward(self, x):
        # 在传播数据时初始化权重
        if self.W.data is None:
            if not x.shape:
                self.in_size = 1
            else:
                self.in_size = x.shape[1]
            self._init_W()

        y = linear(x, self.W, self.b)
        return y
